import os
import sys
import time

import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from PIL import Image
import torch
from models import Discriminator, Generator
from utils import ReplayBuffer, LambdaLR, weights_init_normal
from datasets import ImageDataset
import itertools
import tensorboardX

from torchvision.utils import save_image

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

netG_A2B = torch.load('../model/netG_A2B.pkl')
netG_B2A = torch.load('../model/netG_B2A.pkl')
# print(netG_A2B)
# print(netG_B2A)

netG_A2B.eval()
netG_B2A.eval()

transforms_ = [
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
]

img_size = 256
input_A = torch.ones([1, 3, img_size, img_size], dtype=torch.float).to(device)
input_B = torch.ones([1, 3, img_size, img_size], dtype=torch.float).to(device)
data_root = '../data/'
dataLoader = DataLoader(ImageDataset(data_root, transform=transforms_, model='test'), batch_size=1, shuffle=False,
                        num_workers=8)


def test_model():
    # if not os.path.exists('../data/outputs/A'):
    #     os.mkdir('../data/outputs/A')
    # if not os.path.exists('../data/outputs/B'):
    #     os.mkdir('../data/outputs/B')

    for i, batch in enumerate(dataLoader):
        real_A = torch.tensor(input_A.copy_(batch['A']), dtype=torch.float).to(device)
        real_B = torch.tensor(input_A.copy_(batch['B']), dtype=torch.float).to(device)

        fake_B = 0.5 * (netG_A2B(real_A).data + 1.0)
        fake_A = 0.5 * (netG_B2A(real_B).data + 1.0)

        save_image(fake_A, '../data/outputs/A/{}.png'.format(i))
        save_image(fake_B, '../data/outputs/B/{}.png'.format(i))

        print(i)


if __name__ == '__main__':
    start = time.time()
    test_model()
    end = time.time()
    print(str((end - start)) + 's')
